from qtpy.QtWidgets import *


class FileDialog(QFileDialog):
    """
    自定义文件对话框类，继承自 QFileDialog。
    可根据不同模式设置文件对话框的行为，支持多选功能。

    :param mode: 对话框模式，可选值为 "all", "file", "directory"，默认为 "all"
    :param multifile: 是否启用多选功能，默认为 False
    :param parent: 父控件，默认为 None
    :param kwargs: 传递给父类构造函数的其他关键字参数
    """

    def __init__(self, mode="all", multifile=False, parent=None, **kwargs):
        # 调用父类的构造函数初始化对话框
        super(FileDialog, self).__init__(parent=parent, **kwargs)
        # 设置不使用原生文件对话框
        self.setOption(QFileDialog.DontUseNativeDialog, True)
        if mode == "all":
            # 设置文件模式为选择目录
            self.setFileMode(QFileDialog.Directory)
            # 遍历对话框中的所有 QPushButton 控件
            for but in self.findChildren(QPushButton):
                # 查找文本中包含 "open" 或 "choose" 的按钮
                if "open" in but.text().lower() or "choose" in but.text().lower():
                    # 断开按钮原有的点击信号连接
                    but.clicked.disconnect()
                    # 连接按钮的点击信号到 chooseClicked 方法
                    but.clicked.connect(self.chooseClicked)

        elif mode == "file":
            # 设置文件模式为选择任意文件
            self.setFileMode(QFileDialog.AnyFile)

        elif mode == "directory":
            # 设置文件模式为仅选择目录
            self.setFileMode(QFileDialog.DirectoryOnly)

        if multifile:
            # 查找对话框中的 QListView 控件
            self.listView = self.findChild(QListView)
            if self.listView:
                # 设置 QListView 的选择模式为扩展选择
                self.listView.setSelectionMode(QAbstractItemView.ExtendedSelection)
            # 查找对话框中的 QTreeView 控件
            self.treeView = self.findChild(QTreeView)
            if self.treeView:
                # 设置 QTreeView 的选择模式为扩展选择
                self.treeView.setSelectionMode(QAbstractItemView.ExtendedSelection)

    def chooseClicked(self):
        """
        处理选择按钮点击事件，接受对话框。
        """
        QDialog.accept(self)